from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import openai
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

app = FastAPI()

# 启用CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有源，在生产环境中应该更具体
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 设置OpenAI API密钥
openai.api_key = os.getenv("OPENAI_API_KEY")

# 简单的内存型会话存储
sessions = {}

class ChatRequest(BaseModel):
    session_id: str
    message: str

class ResetRequest(BaseModel):
    session_id: str

@app.post("/chat")
async def chat(request: ChatRequest):
    if request.session_id not in sessions:
        sessions[request.session_id] = []

    sessions[request.session_id].append({"role": "user", "content": request.message})

    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=sessions[request.session_id]
        )
        ai_response = response.choices[0].message['content'].strip()
        sessions[request.session_id].append({"role": "assistant", "content": ai_response})

        return {"response": ai_response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/reset")
async def reset_session(request: ResetRequest):
    if request.session_id in sessions:
        del sessions[request.session_id]
    return {"message": "Session reset successfully"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
